package cn.slipi.common.decorator;

import cn.slipi.common.utils.LogUtils;
import cn.slipi.common.utils.TraceUtil;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;
import reactor.core.scheduler.Schedulers;

import java.util.stream.Collectors;

public class PartnerServerHttpResponseDecorator extends ServerHttpResponseDecorator {


    private static final Logger LOGGER = LoggerFactory.getLogger(PartnerServerHttpResponseDecorator.class);

    PartnerServerHttpResponseDecorator(ServerHttpResponse delegate) {
        super(delegate);
    }

    @Override
    public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
        return super.writeAndFlushWith(body);
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        final String headers = super.getHeaders().entrySet()
                .stream()
                .map(entry -> "" + entry.getKey() + ":" + String.join(";", entry.getValue()))
                .collect(Collectors.joining(";"));
        String traceId = TraceUtil.getCurrentTraceId();
        if (body instanceof Mono) {
            final Mono<DataBuffer> monoBody = (Mono<DataBuffer>) body;
            return super.writeWith(monoBody.publishOn(Schedulers.single()).map(dataBuffer -> {
                TraceUtil.traceStart(traceId);
                return LogUtils.loggingResponse(LOGGER, headers, dataBuffer);
            }).doFinally(this::end));
        } else if (body instanceof Flux) {
            final Flux<DataBuffer> monoBody = (Flux<DataBuffer>) body;
            return super.writeWith(monoBody.publishOn(Schedulers.single()).map(dataBuffer -> {
                TraceUtil.traceStart(traceId);
                return LogUtils.loggingResponse(LOGGER, headers, dataBuffer);
            }).doFinally(this::end));
        }
        return super.writeWith(body);
    }

    private void end(SignalType signalType) {
        TraceUtil.traceEnd();
    }

}